#include <stdio.h>

#include "common/common_types.h"
#include "NanoFlowVis/SubList.h"
#include <spu_mfcio.h>

const int initTag = 0;
const int incomingTag = 1;
const int outgoingTag = 2;

IncomingDataPacket incomingData[SHIFTS_COUNT];
OutgoingDataPacket outgoingData[SHIFTS_COUNT];

SpeInitData initData;
ForcesArray forces;

inline void getIncomingData(unsigned int& packets_count) {
    if (packets_count < initData.packetCount) {
        for (unsigned int shiftIdx = 0; shiftIdx < SHIFTS_COUNT; ++shiftIdx) {
            mfc_get(&incomingData[shiftIdx], initData.incomingDataPtr[shiftIdx], sizeof(IncomingDataPacket), incomingTag, 0, 0);
        }
    }
    ++packets_count;
}

inline void calc() {
    _ZeroStruct(&outgoingData);

    for (unsigned int shift = 0; shift < SHIFTS_COUNT; ++shift) {
        outgoingData[shift].header.accelsCount = incomingData[shift].header.particlesCount;
    }
    for (unsigned int shiftFirst = 0; shiftFirst < SHIFTS_COUNT; ++shiftFirst) {
        for (unsigned int shiftSecond = 0; shiftSecond < SHIFTS_COUNT; ++shiftSecond) {
            if (!shouldInteract(shiftFirst, shiftSecond)) {
                continue;
            }
            const IncomingDataPacket& firstPacket  = incomingData[shiftFirst];
            const IncomingDataPacket& secondPacket = incomingData[shiftSecond];
            OutgoingDataPacket& firstOut = outgoingData[shiftFirst];
            OutgoingDataPacket& secondOut = outgoingData[shiftSecond];
            for (unsigned int firstIdx = 0; firstIdx < firstPacket.header.particlesCount; ++firstIdx) {
                const ParticleData& mainPart = firstPacket.data[firstIdx];
                Vector3& mainAccel = firstOut.accels[firstIdx];
                const Vector3& p = mainPart.pos;
                const unsigned int secondStart = shiftFirst == shiftSecond ? firstIdx + 1 : 0;
                for (unsigned int secondIdx = secondStart; secondIdx < secondPacket.header.particlesCount; ++secondIdx) {
                    const ParticleData& secondPart = secondPacket.data[secondIdx];
                    if (mainPart.moleculaId == secondPart.moleculaId) {
                        continue;
                    }
                    const Vector3 p_c = p - secondPart.pos;
                    const ForceInfo& force = forces.array[mainPart.componentId][secondPart.componentId];
                    //const IFr& a = *Particle::ForceArr[mainPart.componentId][secondPart.componentId];

                    const float dist = p_c.length();
                    if (dist <= 2*force.Sigma) {
                        Vector3& secondAccel = secondOut.accels[secondIdx];

                        const float M1 = mainPart.mass;
                        const float M2 = secondPart.mass;
                        const Vector3 cross_n = p_c / dist;
                        const float Meff = M1 * M2 / (M1 + M2);

                        const float accel = Simple_Calc(force, Meff, dist, initData.tau);

                        const Vector3 v_accel = cross_n * accel;

                        const Vector3 v_accel1 = v_accel * (Meff / M1);
                        const Vector3 v_accel2 = v_accel * (-Meff / M2);

                        mainAccel += v_accel1;
                        secondAccel += v_accel2;
                    }
                }
            }
        }
    }

    //printf("Incoming data: %f\n", incomingData[0].data[0].mass);
}

int main(unsigned long long speid, unsigned long long argp) {
    //printf("Starting...\n");
    // grabbing pointers
    mfc_get(&initData, argp, sizeof(SpeInitData), initTag, 0, 0);
    mfc_write_tag_mask(TAG_MASK(initTag));
    mfc_read_tag_status_all();

    // load forces info
    mfc_get(&forces, initData.forcesInfoPtr, sizeof(ForcesArray), TAG_MASK(incomingTag), 0, 0);
    // main cycle
    unsigned int packets_read = 0;
    while (packets_read < initData.packetCount) {
        // load incoming data packet
        getIncomingData(packets_read);
        // wait for all transfers to complete
        mfc_write_tag_mask(TAG_MASK(incomingTag) | TAG_MASK(outgoingTag));
        mfc_read_tag_status_all();
        // calculate
        calc();
        // push calculated data
        for (unsigned int shiftIdx = 0; shiftIdx < SHIFTS_COUNT; ++shiftIdx) {
            if (outgoingData[shiftIdx].header.accelsCount > 0) {
                mfc_put(&outgoingData[shiftIdx], initData.outgoingDataPtr[shiftIdx], sizeof(OutgoingDataPacket), outgoingTag, 0, 0);
            }
        }
        // update pointers
        memcpy(&initData.incomingDataPtr, &incomingData[0].header.nextIncomingDataPtr, sizeof(initData.incomingDataPtr));
        memcpy(&initData.outgoingDataPtr, &incomingData[0].header.nextOutgoingDataPtr, sizeof(initData.outgoingDataPtr));
    }
    // wait for outgoing transfers to complete
    mfc_write_tag_mask(TAG_MASK(outgoingTag));
    mfc_read_tag_status_all();
    //printf("Done!\n");

    return 0;
}
